#include "HairCore.h"

namespace {
	using namespace Eigen;

	Vector3f gen_middle_particle(const Vector3f start, const Vector3f end) {
		const auto u = end - start;
		const auto v = u.cross(Vector3f::Random());
		return start + 0.5f * u + (XRwy::HairCoreMath::HALF_SQRT_3 * u.norm() / v.norm() * v);
	}

    struct pairhash {
    public:
        template <typename T, typename U>
        std::size_t operator()(const std::pair<T, U> &x) const
        {
            return std::hash<T>()(x.first) ^ std::hash<U>()(x.second);
        }
    };

}

namespace XRwy {
    using namespace std;
    using namespace Eigen;

    NormalMassSpring::NormalMassSpring(HairMassSpringSimulator *simulator_, int idx1_, int idx2_, int strand_idx_, float k_, float l0_):
            simulator(simulator_), idx1(idx1_), idx2(idx2_), k(k_), strand_idx(strand_idx_) {
		//we can provide an initial l0 for init
		if (l0_ <= 0.0f)
        	l0 = (get_pos1() - get_pos2()).norm();
		else
			l0 = l0_;
        k_div_l0 = k / l0;
    }

	NormalMassSpring::NormalMassSpring(HairMassSpringSimulator *simulator_, int idx1_, int idx2_, int strand_idx_, float k_): NormalMassSpring(simulator_, idx1_, idx2_, strand_idx_, k_, -1e20f) {};

    void NormalMassSpring::apply_vector(VectorXf &vec, bool local) const  {
        const auto &pos1 = simulator->pos.segment(idx1 * 3, 3);
        const auto &pos2 = simulator->pos.segment(idx2 * 3, 3);
        Vector3f d = pos2 - pos1;
        float l = d.norm();
        Vector3f force = k_div_l0 * (l  - l0) * (d / l);

        vec.segment((idx1 - local * strand_idx) * 3, 3) += force;
        vec.segment((idx2 - local * strand_idx) * 3, 3) -= force;
    }

    void NormalMassSpring::apply_triplet(vector<Triplet<float>> &triplet, float t, bool local) const  {
        Vector3f d = get_pos2() - get_pos1();
        d.normalize();

        float factor = k_div_l0 * t * t;
        float m[9] = { d(0) * d(0) * factor, d(0) * d(1) * factor, d(0) * d(2) * factor,
                       d(1) * d(0) * factor, d(1) * d(1) * factor, d(1) * d(2) * factor,
                       d(2) * d(0) * factor, d(2) * d(1) * factor, d(2) * d(2) * factor };
        int idx_array[2] = {(idx1 - local * strand_idx) * 3, (idx2 - local * strand_idx) * 3};

        for (const int i1 : idx_array)
            for (const int i2 : idx_array)
                for (int i = 0; i < 3; ++i)
                    for (int j = 0; j < 3; ++j)
                        triplet.emplace_back(i1 + i, i2 + j, (i1 == i2 ? 1.0f : -1.0f) * m[i * 3 + j]);
    }

    NormalMassSpring::SegmentType NormalMassSpring::get_pos(int idx) const { return simulator->pos.segment(idx * 3, 3); }

    NormalMassSpring::SegmentType NormalMassSpring::get_vel(int idx) const { return simulator->vel.segment(idx * 3, 3); }

    NormalMassSpring::SegmentType NormalMassSpring::get_rest_pos(int idx) const { return simulator->rest_pos.segment(idx * 3, 3); }

}

namespace XRwy {
	using namespace std;
	using namespace Eigen;

	HairMassSpringSimulator::HairMassSpringSimulator(const HairReader &hair_reader, const HairMassSpringConfiguration &conf_): conf(conf_) {

		//init the visible_particle_count and visible_particle_count_on_strand
		this->visible_particle_count = hair_reader.get_particle_count();

		//init the particles and add virtual particle
		strand_infos = vector<StrandInfo>(hair_reader.get_strand_count());

		particle_count= 0;
		for (int i = 0; i < hair_reader.get_strand_count(); ++i) {
			auto &strand_info = strand_infos[i];
			const auto &strand = hair_reader.get_strands()[i];

			strand_info.index = particle_count;

			for (int particle_idx = 0; particle_idx < strand.size(); ++particle_idx) {
				strand_info.particle_types.push_back(NORMAL);
				++strand_info.visible_particle_count;
				particle_infos.emplace_back(i, strand[particle_idx]); //the particle is in the "i-th" strand

				if (particle_idx + 1 >= strand.size())
					break;	//if it is the last particle, we don't do any operation, just return

				if ((particle_idx - 1 >= 0 && HairCoreMath::is_colinear(strand[particle_idx - 1], strand[particle_idx], strand[particle_idx + 1]))
				    || (particle_idx + 2 < strand.size() && HairCoreMath::is_colinear(strand[particle_idx], strand[particle_idx + 1], strand[particle_idx + 2]))) {
					//if the line "particle_idx, particle_idx + 1" is colinear with previous or next one
					//gen middle particle
					strand_info.particle_types.push_back(VIRTUAL);
					particle_infos.emplace_back(i, gen_middle_particle(strand[particle_idx], strand[particle_idx + 1]));
				}
			}

			particle_count += strand_info.get_particle_count();
		}

		vector_dim = particle_count * 3;
		rest_pos = VectorXf::Zero(vector_dim);
		vel = VectorXf::Zero(vector_dim);
		for (int i = 0; i < particle_count; ++i)
			for (int j = 0; j < 3; ++j)
				rest_pos[i * 3 + j] = particle_infos[i].rest_pos(j);
		//rest_pos is the pos in the identity matrix
        //the pos now is the rest pos multiply by the initial matrix
        pos = VectorXf::Zero(vector_dim);
        for (int i = 0; i < vector_dim; i += 3) {
            pos.segment<3>(i) = conf.init_transformation * rest_pos.segment<3>(i);
        }
        filter = VectorXf::Constant(vector_dim, 1);
        for (const auto & strand_info : strand_infos)
            filter.segment(strand_info.index * 3, 3) << 0, 0, 0;

        //add spring to it
		int particle_cur = 0;
		for (auto & strand_info : strand_infos) {

			int particle_count = strand_info.get_particle_count();
			const auto & particle_types = strand_info.particle_types;

            strand_info.springs.clear();

			for (int i = 0; i < particle_count; ++i) {
                int global_index = particle_cur + i;
				//add edge spring and extra edge spring with k_edge
				//we assume that extra edge k value is equal to k_edge
				if (i + 1 < particle_count)
					strand_info.springs.push_back(new NormalMassSpring(this, global_index, global_index + 1, strand_info.index, conf.ke)); //we use the local index instead of the global index
				if (i + 2 < particle_count && particle_types[i] == NORMAL && particle_types[i + 2] == NORMAL && particle_types[i + 1] == VIRTUAL)
					strand_info.springs.push_back(new NormalMassSpring(this, global_index, global_index + 2, strand_info.index, conf.ke));

				//add bending springs
				int i_next = i + 2;
				//the bending spring always link particles with the same particle types
				while (i_next < particle_count && particle_types[i] != particle_types[i_next])
					++i_next;
				if (i_next < particle_count)
					strand_info.springs.push_back(new NormalMassSpring(this, global_index, global_index - i + i_next, strand_info.index, conf.kb));

				//add torsion springs
				if (i + 3 < particle_count)
					strand_info.springs.push_back(new NormalMassSpring(this, global_index, global_index + 3, strand_info.index, conf.kt));
			}

			particle_cur += strand_info.get_particle_count();
		}

        //add hair self-replusion spring
        aligned_x_particle_indexes.reserve(particle_count);
        for (int i = 0; i < particle_count; ++i)
            aligned_x_particle_indexes.push_back(i);

        stiction_springs.clear();
        update_stiction_spring();
	}

	void HairMassSpringSimulator::update_aligned_hierarchy() {
		const auto & compare_pos = pos;
		std::sort(aligned_x_particle_indexes.begin(), aligned_x_particle_indexes.end(), [&compare_pos](int lhs, int rhs) -> bool {
			return compare_pos(lhs * 3) < compare_pos(rhs * 3);
		});
	}

	int HairMassSpringSimulator::get_particle_count() const {
		return visible_particle_count;
	}

	int HairMassSpringSimulator::get_strand_count() const {
		return strand_infos.size();
	}

	int HairMassSpringSimulator::get_particle_count_on_strand(int idx) const {
		return strand_infos[idx].get_visible_particle_count();
	}

	void HairMassSpringSimulator::copy_particles(float *positions) const {
		int positions_cur = 0;
		for (const auto & strand_info : strand_infos) {
			for (int i = 0; i < strand_info.get_particle_count(); ++i)
				if (strand_info.particle_types[i] == NORMAL) {
					const auto & p = this->pos.segment((strand_info.index + i) * 3, 3);
					positions[positions_cur] = p(0);
					positions[positions_cur + 1] = p(1);
					positions[positions_cur + 2] = p(2);
					positions_cur += 3;
				}
		}
	}

	void HairMassSpringSimulator::on_frame(const HairSimulatorData &data) {
		float t = data.delta_time /conf.npass;

        vector<Vector3f> strand_root_velocities;
        for (const auto & strand : this->strand_infos) {
            Vector3f old_pos = pos.segment<3>(strand.index * 3);
            Vector3f new_pos = data.motion_transform * rest_pos.segment<3>(strand.index * 3);
            Vector3f vel = (new_pos - old_pos) / data.delta_time;
            strand_root_velocities.push_back(vel);
        }

		data.sdf_ptr->update(data.sdf_world2local, data.delta_time);

		for (int n = 0; n < conf.npass; ++n) {
			resolve_time_integration(data, strand_root_velocities, t);

            if (conf.enable_hair_stiction)
			    resolve_stiction(data, t);

            for (int i = 0; i < strand_infos.size(); ++i)
                vel.segment(strand_infos[i].index * 3, 3) << strand_root_velocities[i];
            pos += t * vel;

            if (conf.enable_body_collision)
                resolve_collision(data, t);
        }
	}

	void HairMassSpringSimulator::resolve_time_integration(const HairSimulatorData &data, const vector<Vector3f> strand_root_vels, float t) {
        tbb::parallel_for(size_t(0), strand_infos.size(), [this, t](size_t i){
            auto & strand = strand_infos[i];
            int dim = strand.get_particle_count() * 3;

            VectorXf forces = VectorXf::Zero(dim);
            vector<Triplet<float>> triplets;
            for (const auto & spring : strand.springs) {
                spring->apply_vector(forces, true);
                spring->apply_triplet(triplets, t, true);
            }

            for (int i = 0; i < dim; ++i)
                triplets.emplace_back(i, i, conf.mass + conf.b * t);

            SparseMatrix<float> B(dim, dim);
            B.setFromTriplets(triplets.begin(), triplets.end());

            for (int i = 0; i < dim; i += 3)
                forces.segment<3>(i) += conf.extra_force;

            ConjugateGradient<SparseMatrix<float>, Lower | Upper> cg;
            cg.compute(B);

            int vector_begin_index = strand.index * 3;
            VectorXf tmp_vel = cg.solve(t * forces + conf.mass * vel.segment(vector_begin_index, dim));
            vel.segment(vector_begin_index, dim) << tmp_vel;
        });
	}

    void HairMassSpringSimulator::update_stiction_spring() {
        //Remove some spring that has the distance larger than the invalid dr
        vector<NormalMassSpring> new_stiction_springs;
        vector<int> linked_counts(particle_count, 0);
        unordered_set<pair<int, int>, pairhash> links;
        new_stiction_springs.reserve(particle_count * conf.max_replusion_link);

        for (const auto &spring : stiction_springs)
            if ((spring.get_pos1() - spring.get_pos2()).norm() < conf.invalid_dr) {
                new_stiction_springs.push_back(spring);
                ++linked_counts[spring.idx1];
                ++linked_counts[spring.idx2];
                links.insert(make_pair(spring.idx1, spring.idx2));
                links.insert(make_pair(spring.idx2, spring.idx1));
            }

        stiction_springs = std::move(new_stiction_springs);

        update_aligned_hierarchy();

        for (int i = 0; i < this->particle_count; ++i) {
            int index_i = aligned_x_particle_indexes[i];
            Vector3f pos_i = pos.segment<3>(index_i * 3);

            for (int j = i + 1; linked_counts[index_i] < conf.max_replusion_link && j < this->particle_count; ++j) {
                int index_j = aligned_x_particle_indexes[j];
                Vector3f pos_j = pos.segment<3>(index_j * 3);

                if (pos_j(0) - pos_i(0) > conf.dr)
                    break;

                float dist = (pos_i - pos_j).norm();
                if (dist < conf.dr
                    && particle_infos[index_i].strand_index != particle_infos[index_j].strand_index
                    && linked_counts[index_j] < conf.max_replusion_link
                    && links.find(make_pair(index_i, index_j)) == links.end()) {
                    ++linked_counts[index_i];
                    ++linked_counts[index_j];
                    //0 stands for no strand index
                    stiction_springs.emplace_back(this, index_i, index_j, 0, conf.kr, conf.dr);
                }
            }
        }

        cout << "Total self-replusion link count is: " << stiction_springs.size() / 2 << endl;
        cout << "Average link count is: " << (double)(stiction_springs.size() / 2) / this->particle_count << endl;
    }

    void HairMassSpringSimulator::resolve_stiction(const HairSimulatorData &data, float t) {
        update_stiction_spring();

        VectorXf forces = VectorXf::Zero(this->vector_dim);
//		vector<Triplet<float>> triplets;
//		SparseMatrix<float> B(this->vector_dim, this->vector_dim);
//
//		for (const auto &spring : stiction_springs) {
//			spring.apply_vector(forces, false); //use global index
//			spring.apply_triplet(triplets, t, false); //use global index
//		}
//
//		for (int i = 0; i < vector_dim; ++i)
//			triplets.emplace_back(i, i, conf.mass);
//
//		B.setFromTriplets(triplets.begin(), triplets.end());
//
//		ConjugateGradient<SparseMatrix<float>, Lower | Upper> cg;
//		cg.compute(B);
//		VectorXf tmp_vel = cg.solve(t * forces + conf.mass * vel);
//        vel = std::move(tmp_vel);

        for (const auto &spring : stiction_springs)
            spring.apply_vector(forces, false);
        vel += (t / conf.mass) * forces;
    }

    void HairMassSpringSimulator::resolve_collision(const HairSimulatorData &data, float t) {
        //data.sdf_ptr->fix(pos, vel, filter, t);
        for (size_t i = 0; i < pos.rows(); i += 3)
            if (filter(i) != 0) {
				Vector3f tmp_pos = pos.segment<3>(i);
				Vector3f tmp_vel = vel.segment<3>(i);
                data.sdf_ptr->fix(tmp_pos, tmp_vel, t);
				pos.segment<3>(i) = tmp_pos;
				vel.segment<3>(i) = tmp_vel;
            }
    }

	vector<pair<HairParticleType, Vector3f>> HairMassSpringSimulator::dump_all_particles_info() const {
		vector<pair<HairParticleType, Vector3f>> ret;
		ret.reserve(particle_count);

		int cur = 0;
		for (const auto & strand_info : strand_infos) {
			for (const auto & particle_type : strand_info.particle_types) {
				Vector3f tmp;
				tmp << pos(cur), pos(cur + 1), pos(cur + 2);
				cur += 3;
				ret.push_back(make_pair(particle_type, tmp));
			}
		}

		return std::move(ret);
	}

	HairMassSpringSimulator::~HairMassSpringSimulator() {
		for (auto &strand : strand_infos)
			for (auto *spring : strand.springs)
				delete spring;
	}
}
